#include <Eigen/Core>
#include <ceres/ceres.h>
#include <ceres/rotation.h>

#include <vector>

#include <autodiff/reverse.hpp>
#include <autodiff/reverse/eigen.hpp>

using namespace Eigen;

// 相平面坐标
std::vector<Vector2d> point2ds = {Vector2d(-86.15, -68.99), Vector2d(-53.40, 82.21), Vector2d(-14.78, -76.63),
                                  Vector2d(10.46, 64.43)};
// 控制点坐标
std::vector<Vector3d> point3ds = {Vector3d(36589.41, 25273.32, 2195.17), Vector3d(37631.08, 31324.51, 728.69),
                                  Vector3d(39100.97, 24934.98, 2386.50), Vector3d(40426.54, 30319.81, 757.31)};
// 内方位元素
double f = 153.24;
double x_0 = 0.0;
double y_0 = 0.0;

using namespace autodiff;

Matrix3var rodrigues(const Vector3var &angle_axis) {
    static const double kOne = 1.0;
    const var theta2 = angle_axis.squaredNorm();
    Matrix3var R;
    if (theta2 > var(std::numeric_limits<double>::epsilon())) {
        // We want to be careful to only evaluate the square root if the
        // norm of the angle_axis vector is greater than zero. Otherwise
        // we get a division by zero.
        const var theta = sqrt(theta2);
        const var wx = angle_axis[0] / theta;
        const var wy = angle_axis[1] / theta;
        const var wz = angle_axis[2] / theta;

        const var costheta = cos(theta);
        const var sintheta = sin(theta);

        R(0, 0) = costheta + wx * wx * (kOne - costheta);
        R(1, 0) = wz * sintheta + wx * wy * (kOne - costheta);
        R(2, 0) = -wy * sintheta + wx * wz * (kOne - costheta);
        R(0, 1) = wx * wy * (kOne - costheta) - wz * sintheta;
        R(1, 1) = costheta + wy * wy * (kOne - costheta);
        R(2, 1) = wx * sintheta + wy * wz * (kOne - costheta);
        R(0, 2) = wy * sintheta + wx * wz * (kOne - costheta);
        R(1, 2) = -wx * sintheta + wy * wz * (kOne - costheta);
        R(2, 2) = costheta + wz * wz * (kOne - costheta);
    } else {
        // Near zero, we switch to using the first order Taylor expansion.
        R(0, 0) = kOne;
        R(1, 0) = angle_axis[2];
        R(2, 0) = -angle_axis[1];
        R(0, 1) = -angle_axis[2];
        R(1, 1) = kOne;
        R(2, 1) = angle_axis[0];
        R(0, 2) = angle_axis[1];
        R(1, 2) = -angle_axis[0];
        R(2, 2) = kOne;
    }
    return R;
}

std::tuple<Vector3d, Matrix3d, double, Vector3d> question4() {
    Vector3d Xs = Vector3d::Zero();
    Matrix3d R = Matrix3d::Identity();
    double sigma_0 = 0.0;
    Vector3d m_Xs = Vector3d::Zero();

    // TODO 实现后方交会

    // 1 计算外方位线元素初值

    // 2 角元素初值就为单位阵

    // 3 迭代优化求解

    return std::make_tuple(Xs, R, sigma_0, m_Xs);
}

struct ProjectionResidual {
    ProjectionResidual(const Vector2d &point2d, const Vector3d &point3d) : point2d(point2d), point3d(point3d) {}

    template <typename T> bool operator()(const T *Xs, const T *r, T *residual) const {
        T R[9];
        ceres::AngleAxisToRotationMatrix(r, ceres::RowMajorAdapter3x3(R));

        T X = point3d.x() - Xs[0];
        T Y = point3d.y() - Xs[1];
        T Z = point3d.z() - Xs[2];

        T x = -f * (R[0] * X + R[3] * Y + R[6] * Z) / (R[2] * X + R[5] * Y + R[8] * Z) + x_0;
        T y = -f * (R[1] * X + R[4] * Y + R[7] * Z) / (R[2] * X + R[5] * Y + R[8] * Z) + y_0;

        residual[0] = point2d.x() - x;
        residual[1] = point2d.y() - y;

        return true;
    }

private:
    Vector2d point2d;
    Vector3d point3d;
};

std::tuple<Vector3d, Matrix3d, double, Vector3d> golden4() {
    Vector3d Xs = Vector3d::Zero();
    Matrix3d R = Matrix3d::Identity();
    double sigma_0 = 0.0;
    Vector3d m_Xs = Vector3d::Zero();

    // 计算初值
    for (const auto &point : point3ds) {
        Xs += point;
    }
    Xs = Xs / point3ds.size();
    Xs.z() = 8000.0;

    // 角元素初值就是单位阵
    auto angle_axis = Eigen::AngleAxisd(R);
    Vector3d r = angle_axis.angle() * angle_axis.axis();

    ceres::Problem problem;

    std::vector<ceres::ResidualBlockId> blocks;

    for (int i = 0; i < point3ds.size(); ++i) {
        ceres::CostFunction *cost = new ceres::AutoDiffCostFunction<ProjectionResidual, 2, 3, 3>(
            new ProjectionResidual(point2ds[i], point3ds[i]));
        ceres::ResidualBlockId block = problem.AddResidualBlock(cost, nullptr, Xs.data(), r.data());
        blocks.push_back(block);
    }

    ceres::Solver::Options op;
    op.max_num_iterations = 100;
    op.linear_solver_type = ceres::DENSE_QR;

    ceres::Solver::Summary summary;
    ceres::Solve(op, &problem, &summary);

    // 计算单位权中误差
    {
        ceres::Problem::EvaluateOptions op_eval;
        op_eval.residual_blocks = blocks;
        double cost;
        problem.Evaluate(op_eval, &cost, nullptr, nullptr, nullptr);
        cost *= 2.0; // ceres 记录的是包含 1/2 的
        sigma_0 = std::sqrt(cost / (2 * point2ds.size() - 6));
    }

    // 计算协方差矩阵
    {
        Matrix3d cov_Xs;
        std::vector<std::pair<const double *, const double *>> covariance_blocks;
        covariance_blocks.emplace_back(Xs.data(), Xs.data());
        ceres::Covariance::Options options;
        ceres::Covariance covariance(options);
        covariance.Compute(covariance_blocks, &problem);
        covariance.GetCovarianceBlock(Xs.data(), Xs.data(), cov_Xs.data());

        m_Xs.x() = sigma_0 * std::sqrt(cov_Xs(0, 0));
        m_Xs.y() = sigma_0 * std::sqrt(cov_Xs(1, 1));
        m_Xs.z() = sigma_0 * std::sqrt(cov_Xs(2, 2));
    }

    // 计算旋转矩阵
    {
        angle_axis.angle() = r.norm();
        angle_axis.axis() = r.normalized();

        R = angle_axis.toRotationMatrix();
    }
    return std::make_tuple(Xs, R, sigma_0, m_Xs);
}
